#include <mpi.h>
#include <stdio.h>
#include "utils.h"

int main(int argc, char** argv) {
    int rank, size;
    int data;
    MPI_Comm local_comm, peer_comm;

    // 初始化 MPI 环境
    MPI_Init(&argc, &argv);

    // 获取当前进程的 rank 和总的进程数
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    // 创建两个组，每个组包含一半的进程
    int color = rank < size / 2 ? 0 : 1;
    // INFO("rank = %d, color = %d\n", rank, color);
    MPI_Comm_split(MPI_COMM_WORLD, color, rank, &local_comm);

    // 只有第一个组的 rank 0 初始化数据
    if (rank == 0) {
        data = 42;
        // INFO("Group 0, Process 0 broadcasting data %d\n", data);
    }


    if (color == 0) {
        // 2表示peer在MPI_COMM_WORLD中的rank
        MPI_Intercomm_create(local_comm, 0, MPI_COMM_WORLD, 2, 100, &peer_comm);
    } else {
        MPI_Intercomm_create(local_comm, 0, MPI_COMM_WORLD, 0, 100, &peer_comm);
    }
    MPI_Barrier(MPI_COMM_WORLD);

    if (color == 0) {
        if (rank == 0)
            MPI_Bcast(&data, 1, MPI_INT, MPI_ROOT, peer_comm);
        else
           MPI_Bcast(&data, 1, MPI_INT, MPI_PROC_NULL, peer_comm);
    } else {
        //! 这里极其重要，不然结果会不正确
        MPI_Bcast(&data, 1, MPI_INT, 0, peer_comm);
    }
    MPI_Barrier(MPI_COMM_WORLD);

    // 所有进程打印接收到的数据
    INFO("Group %d, Process %d received data %d\n", color, rank, data);

    // 释放组间通信子
    MPI_Comm_free(&peer_comm);
    MPI_Comm_free(&local_comm);

    // 结束 MPI 环境
    MPI_Finalize();

    return 0;
}